{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import numpy as np\n",
    "import matplotlib.dates as mdates\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import norm\n",
    "import numdifftools as nd\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dataset = []\n",
    "\n",
    "# read the measurement data, use 0.0 to stand LIDAR data\n",
    "# and 1.0 stand RADAR data\n",
    "with open('data_synthetic.txt', 'r') as f:\n",
    "    lines = f.readlines()\n",
    "    for line in lines:\n",
    "        line = line.strip('\\n')\n",
    "        line = line.strip()\n",
    "        numbers = line.split()\n",
    "        result = []\n",
    "        for i, item in enumerate(numbers):\n",
    "            item.strip()\n",
    "            if i == 0:\n",
    "                if item == 'L':\n",
    "                    result.append(0.0)\n",
    "                else:\n",
    "                    result.append(1.0)\n",
    "            else:\n",
    "                result.append(float(item))\n",
    "        dataset.append(result)\n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1. 0. 0. 0. 0.]\n",
      " [0. 1. 0. 0. 0.]\n",
      " [0. 0. 1. 0. 0.]\n",
      " [0. 0. 0. 1. 0.]\n",
      " [0. 0. 0. 0. 1.]] (5, 5)\n",
      "[[1. 0. 0. 0. 0.]\n",
      " [0. 1. 0. 0. 0.]] (2, 5)\n",
      "[[0.0225 0.    ]\n",
      " [0.     0.0225]] (2, 2)\n",
      "[[0.09   0.     0.    ]\n",
      " [0.     0.0009 0.    ]\n",
      " [0.     0.     0.09  ]] (3, 3)\n"
     ]
    }
   ],
   "source": [
    "P = np.diag([1.0, 1.0, 1.0, 1.0, 1.0])\n",
    "print(P, P.shape)\n",
    "H_lidar = np.array([[ 1.,  0.,  0.,  0.,  0.],\n",
    "       [ 0.,  1.,  0.,  0.,  0.]])\n",
    "print(H_lidar, H_lidar.shape)\n",
    "\n",
    "R_lidar = np.array([[0.0225, 0.],[0., 0.0225]])\n",
    "R_radar = np.array([[0.09, 0., 0.],[0., 0.0009, 0.], [0., 0., 0.09]])\n",
    "print(R_lidar, R_lidar.shape)\n",
    "print(R_radar, R_radar.shape)\n",
    "\n",
    "# process noise standard deviation for a\n",
    "std_noise_a = 2.0\n",
    "# process noise standard deviation for yaw acceleration\n",
    "std_noise_yaw_dd = 0.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def control_psi(psi):\n",
    "    while (psi > np.pi or psi < -np.pi):\n",
    "        if psi > np.pi:\n",
    "            psi = psi - 2 * np.pi\n",
    "        if psi < -np.pi:\n",
    "            psi = psi + 2 * np.pi\n",
    "    return psi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialize with LIDAR measurement!\n",
      "[0.3122427 0.5803398 0.        0.        0.       ] (5,)\n"
     ]
    }
   ],
   "source": [
    "state = np.zeros(5)\n",
    "init_measurement = dataset[0]\n",
    "current_time = 0.0\n",
    "if init_measurement[0] == 0.0:\n",
    "    print('Initialize with LIDAR measurement!')\n",
    "    current_time = init_measurement[3]\n",
    "    state[0] = init_measurement[1]\n",
    "    state[1] = init_measurement[2]\n",
    "\n",
    "else:\n",
    "    print('Initialize with RADAR measurement!')\n",
    "    current_time = init_measurement[4]\n",
    "    init_rho = init_measurement[1]\n",
    "    init_psi = init_measurement[2]\n",
    "    init_psi = control_psi(init_psi)\n",
    "    state[0] = init_rho * np.cos(init_psi)\n",
    "    state[1] = init_rho * np.sin(init_psi)\n",
    "print(state, state.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Preallocation for Saving\n",
    "px = []\n",
    "py = []\n",
    "vx = []\n",
    "vy = []\n",
    "\n",
    "gpx = []\n",
    "gpy = []\n",
    "gvx = []\n",
    "gvy = []\n",
    "\n",
    "mx = []\n",
    "my = []\n",
    "\n",
    "def savestates(ss, gx, gy, gv1, gv2, m1, m2):\n",
    "    px.append(ss[0])\n",
    "    py.append(ss[1])\n",
    "    vx.append(np.cos(ss[3]) * ss[2])\n",
    "    vy.append(np.sin(ss[3]) * ss[2])\n",
    "\n",
    "    gpx.append(gx)\n",
    "    gpy.append(gy)\n",
    "    gvx.append(gv1)\n",
    "    gvy.append(gv2)\n",
    "    mx.append(m1)\n",
    "    my.append(m2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 状态转移函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "measurement_step = len(dataset)\n",
    "state = state.reshape([5, 1])\n",
    "dt = 0.05\n",
    "\n",
    "I = np.eye(5)\n",
    "\n",
    "transition_function = lambda y: np.vstack((\n",
    "    y[0] + (y[2] / y[4]) * (np.sin(y[3] + y[4] * dt) - np.sin(y[3])),\n",
    "    y[1] + (y[2] / y[4]) * (-np.cos(y[3] + y[4] * dt) + np.cos(y[3])),\n",
    "    y[2],\n",
    "    y[3] + y[4] * dt,\n",
    "    y[4]))\n",
    "\n",
    "# when omega is 0\n",
    "transition_function_1 = lambda m: np.vstack((m[0] + m[2] * np.cos(m[3]) * dt,\n",
    "                                             m[1] + m[2] * np.sin(m[3]) * dt,\n",
    "                                             m[2],\n",
    "                                             m[3] + m[4] * dt,\n",
    "                                             m[4]))\n",
    "\n",
    "J_A = nd.Jacobian(transition_function)\n",
    "J_A_1 = nd.Jacobian(transition_function_1)\n",
    "# print(J_A([1., 2., 3., 4., 5.]))\n",
    "\n",
    "measurement_function = lambda k: np.vstack((np.sqrt(k[0] * k[0] + k[1] * k[1]),\n",
    "                                            math.atan2(k[1], k[0]),\n",
    "                                            (k[0] * k[2] * np.cos(k[3]) + k[1] * k[2] * np.sin(k[3])) / np.sqrt(k[0] * k[0] + k[1] * k[1])))\n",
    "J_H = nd.Jacobian(measurement_function)\n",
    "# J_H([1., 2., 3., 4., 5.])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for step in range(1, measurement_step):\n",
    "\n",
    "    # Prediction\n",
    "    # ========================\n",
    "    t_measurement = dataset[step]\n",
    "    if t_measurement[0] == 0.0:\n",
    "        m_x = t_measurement[1]\n",
    "        m_y = t_measurement[2]\n",
    "        z = np.array([[m_x], [m_y]])\n",
    "\n",
    "        dt = (t_measurement[3] - current_time) / 1000000.0\n",
    "        current_time = t_measurement[3]\n",
    "\n",
    "        # true position\n",
    "        g_x = t_measurement[4]\n",
    "        g_y = t_measurement[5]\n",
    "        g_v_x = t_measurement[6]\n",
    "        g_v_y = t_measurement[7]\n",
    "\n",
    "    else:\n",
    "        m_rho = t_measurement[1]\n",
    "        m_psi = t_measurement[2]\n",
    "        m_dot_rho = t_measurement[3]\n",
    "        z = np.array([[m_rho], [m_psi], [m_dot_rho]])\n",
    "\n",
    "        dt = (t_measurement[4] - current_time) / 1000000.0\n",
    "        current_time = t_measurement[4]\n",
    "\n",
    "        # true position\n",
    "        g_x = t_measurement[5]\n",
    "        g_y = t_measurement[6]\n",
    "        g_v_x = t_measurement[7]\n",
    "        g_v_y = t_measurement[8]\n",
    "\n",
    "    if np.abs(state[4, 0]) < 0.0001:  # omega is 0, Driving straight\n",
    "        state = transition_function_1(state.ravel().tolist())\n",
    "        state[3, 0] = control_psi(state[3, 0])\n",
    "        JA = J_A_1(state.ravel().tolist())\n",
    "    else:  # otherwise\n",
    "        state = transition_function(state.ravel().tolist())\n",
    "        state[3, 0] = control_psi(state[3, 0])\n",
    "        JA = J_A(state.ravel().tolist())\n",
    "\n",
    "\n",
    "    G = np.zeros([5, 2])\n",
    "    G[0, 0] = 0.5 * dt * dt * np.cos(state[3, 0])\n",
    "    G[1, 0] = 0.5 * dt * dt * np.sin(state[3, 0])\n",
    "    G[2, 0] = dt\n",
    "    G[3, 1] = 0.5 * dt * dt\n",
    "    G[4, 1] = dt\n",
    "\n",
    "    Q_v = np.diag([std_noise_a*std_noise_a, std_noise_yaw_dd*std_noise_yaw_dd])\n",
    "    Q = np.dot(np.dot(G, Q_v), G.T)\n",
    "\n",
    "    # Project the error covariance ahead\n",
    "    P = np.dot(np.dot(JA, P), JA.T) + Q\n",
    "\n",
    "    # Measurement Update (Correction)\n",
    "    # ===============================\n",
    "    if t_measurement[0] == 0.0:\n",
    "        # Lidar\n",
    "        S = np.dot(np.dot(H_lidar, P), H_lidar.T) + R_lidar\n",
    "        K = np.dot(np.dot(P, H_lidar.T), np.linalg.inv(S))\n",
    "\n",
    "        y = z - np.dot(H_lidar, state)\n",
    "        y[1, 0] = control_psi(y[1, 0])\n",
    "        state = state + np.dot(K, y)\n",
    "        state[3, 0] = control_psi(state[3, 0])\n",
    "        # Update the error covariance\n",
    "        P = np.dot((I - np.dot(K, H_lidar)), P)\n",
    "\n",
    "        # Save states for Plotting\n",
    "        savestates(state.ravel().tolist(), g_x, g_y, g_v_x, g_v_y, m_x, m_y)\n",
    "\n",
    "    else:\n",
    "        # Radar\n",
    "        JH = J_H(state.ravel().tolist())\n",
    "\n",
    "        S = np.dot(np.dot(JH, P), JH.T) + R_radar\n",
    "        K = np.dot(np.dot(P, JH.T), np.linalg.inv(S))\n",
    "        map_pred = measurement_function(state.ravel().tolist())\n",
    "        if np.abs(map_pred[0, 0]) < 0.0001:\n",
    "            # if rho is 0\n",
    "            map_pred[2, 0] = 0\n",
    "\n",
    "        y = z - map_pred\n",
    "        y[1, 0] = control_psi(y[1, 0])\n",
    "\n",
    "        state = state + np.dot(K, y)\n",
    "        state[3, 0] = control_psi(state[3, 0])\n",
    "        # Update the error covariance\n",
    "        P = np.dot((I - np.dot(K, JH)), P)\n",
    "\n",
    "        savestates(state.ravel().tolist(), g_x, g_y, g_v_x, g_v_y, m_rho * np.cos(m_psi), m_rho * np.sin(m_psi))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.07363360908934428 0.08045989331938738 0.22916598526423182 0.3099938876610046\n"
     ]
    }
   ],
   "source": [
    "def rmse(estimates, actual):\n",
    "    result = np.sqrt(np.mean((estimates-actual)**2))\n",
    "    return result\n",
    "\n",
    "print(rmse(np.array(px), np.array(gpx)),\n",
    "      rmse(np.array(py), np.array(gpy)),\n",
    "      rmse(np.array(vx), np.array(gvx)),\n",
    "      rmse(np.array(vy), np.array(gvy)))\n",
    "\n",
    "# write to the output file\n",
    "stack = [px, py, vx, vy, mx, my, gpx, gpy, gvx, gvy]\n",
    "stack = np.array(stack)\n",
    "stack = stack.T\n",
    "np.savetxt('output.csv', stack, '%.6f')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 绘制结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub data rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_data_rate_limit`.\n",
      "/usr/local/lib/python3.5/dist-packages/plotly/offline/offline.py:459: UserWarning:\n",
      "\n",
      "Your filename `EKF` didn't end with .html. Adding .html to the end of your file.\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'file:///devdisk/tmp/EKF.html'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import plotly.offline as py\n",
    "from plotly.graph_objs import *\n",
    "import pandas as pd\n",
    "import math\n",
    "\n",
    "py.init_notebook_mode()\n",
    "\n",
    "my_cols = ['px_est', 'py_est', 'vx_est', 'vy_est', 'px_meas', 'py_meas', 'px_gt', 'py_gt', 'vx_gt', 'vy_gt']\n",
    "with open('output.csv') as f:\n",
    "    table_ekf_output = pd.read_table(f, sep=' ', header=None, names=my_cols, lineterminator='\\n')\n",
    "\n",
    "    # table_ekf_output\n",
    "\n",
    "import plotly.offline as py\n",
    "from plotly.graph_objs import *\n",
    "\n",
    "# Measurements\n",
    "trace2 = Scatter(\n",
    "    x=table_ekf_output['px_meas'],\n",
    "    y=table_ekf_output['py_meas'],\n",
    "    xaxis='x2',\n",
    "    yaxis='y2',\n",
    "    name='Measurements',\n",
    "    mode = 'markers'\n",
    ")\n",
    "\n",
    "# estimations\n",
    "trace1 = Scatter(\n",
    "    x=table_ekf_output['px_est'],\n",
    "    y=table_ekf_output['py_est'],\n",
    "    xaxis='x2',\n",
    "    yaxis='y2',\n",
    "    name='KF- Estimate',\n",
    "    mode='markers'\n",
    ")\n",
    "\n",
    "# Ground Truth\n",
    "trace3 = Scatter(\n",
    "    x=table_ekf_output['px_gt'],\n",
    "    y=table_ekf_output['py_gt'],\n",
    "    xaxis='x2',\n",
    "    yaxis='y2',\n",
    "    name='Ground Truth',\n",
    "    mode='markers'\n",
    ")\n",
    "\n",
    "data = [trace1, trace2, trace3]\n",
    "\n",
    "layout = Layout(\n",
    "    xaxis2=dict(\n",
    "\n",
    "        anchor='x2',\n",
    "        title='px'\n",
    "    ),\n",
    "    yaxis2=dict(\n",
    "\n",
    "        anchor='y2',\n",
    "        title='py'\n",
    "    )\n",
    ")\n",
    "\n",
    "fig = Figure(data=data, layout=layout)\n",
    "py.plot(fig, filename='EKF')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
